import io
import socket
import os
import json
import PySimpleGUI as sg
import threading
import yaml
from logging.config import dictConfig
import logging
import paho.mqtt.client as mqtt
import serial
import struct
import time
import datetime
LOCAL_TIMEZONE = datetime.datetime.now(
    datetime.timezone.utc).astimezone().tzinfo
# MQTT接口

# YAML文件读写


logging.basicConfig(level=logging.INFO)
# 计算程序文件所在的实际目录，文件名和无扩展的文件名
real_path = os.path.realpath(__file__)
logging.info('real_path = {}'.format(real_path))
base_name = os.path.basename(real_path)
logging.info('base_name = {}'.format(base_name))
base_dir = os.path.dirname(real_path)
logging.info('base_dir = {}'.format(base_dir))
os.chdir(base_dir)
logging.info('change dir to base_dir = {}'.format(os.getcwd()))
program_name, _ext = os.path.splitext(base_name)
init_config_file = os.path.join(base_dir, program_name+'.yaml')  # 计算配置文件
logging.info('init_config_file = {}'.format(init_config_file))

delay_time = 15
w = sg.Window(
    f'get config file, timeout in {delay_time} seconds.',
    [
        [sg.Text('config filename')],
        [
            sg.Input(init_config_file, key='config_file'),
            sg.FileBrowse(
                initial_folder=base_dir,
                file_types=(("YAML configuration file", "*.yaml"),))],
        [sg.OK(), sg.Cancel()]
    ], finalize=True)
event, values = w.read(close=True, timeout=delay_time*1000)

if event in (sg.WIN_CLOSED, 'Cancel'):
    logging.info(' or user closed the window or user cancel, do not continue')
    exit()
elif event in (sg.TIMEOUT_KEY,):
    print(event, values)
    config_file = init_config_file
    logging.info(f'time out, use init config file = {config_file}')
else:
    config_file = values['config_file']
    logging.info(f'user choose config_file = {config_file}')

with open(config_file, 'r', encoding='utf-8') as config_file_io:
    config = yaml.load(config_file_io.read(), yaml.Loader)
    logging.info('config:')
    logging.info(config)
    dictConfig(config['logging'])

logger = logging.getLogger()
udp_daq = logging.getLogger('udp_daq')
serial_daq = logging.getLogger('serial_daq')


class UDP_DAQ():
    """执行UDP数据采集的类"""
    udp_names = 'time_boot_us,lat,lon,alt,homelat,homelon,homealt,vx,vz,distance,vol,roll,pitch,yaw'.split(
        ',')
    udp_units = 'us,deg,deg,m,deg,deg,m,m/s,m/s,m,V,rad,rad,rad'.split(',')
    udp_scalings = [1, 1e7, 1e7, 1e3, 1e7,
                    1e7, 1e3, 1e2, 1e2, 1e2, 1e3, 1, 1, 1]
    udp_offsets = [0]*len(udp_names)
    udp_msg_format = '<IiiiiiihhhHfff'
    source_port = 15550

    wind_speed_names = 'wX,wY,wZ,wS,wD'.split(',')
    wind_speed_units = 'm/s,m/s,m/s,m/s,degree'.split(',')

    def __init__(self, data_logger, mqttc, udp_ip='', udp_port=15550, source_port=15550, ser_daq_obj=None) -> None:
        """初始化"""

        # 初始化变量
        self.data_logger = data_logger
        self.source_port = source_port
        self.ser_daq_obj = ser_daq_obj
        self.vx = 0
        self.vz = 0
        self.roll = 0
        self.pitch = 0
        self.yaw = 0
        self.mqttc = mqttc
        self.client_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.client_socket.setblocking(True)
        # client_socket.settimeout(1.0)
        self.client_socket.bind((udp_ip, udp_port))

        # 写入数据文件的首行
        sio = io.StringIO()
        print('#time\ttimestamp\t'+'\t'.join(self.udp_names), file=sio)
        print('#rfc3339\tseconds\t'+'\t'.join(self.udp_units), file=sio)
        header_rows = sio.getvalue().strip()
        sio.close()
        data_logger.info(header_rows)

        # thingsboard的设备client-side attribute中写入单位信息
        units_client_attribute = {'units': dict(
            zip(self.udp_names+self.wind_speed_names, self.udp_units+self.wind_speed_units))}
        logger.info(f' udp units_client_attribute = {units_client_attribute}')
        payload = json.dumps(units_client_attribute)
        (rc, mid) = mqttc.publish(
            topic='v1/devices/me/attributes',
            payload=payload,
            qos=1)

        # 数采线程控制
        self.event = threading.Event()
        self.running_thread = threading.Thread(
            target=self.daq_func, daemon=True, name='udp_daq')

        # 用sympy推导绝对风速的计算公式
        self.deduce_wind_speed()

    def deduce_wind_speed(self):
        # 用sympy推导绝对风速的计算公式
        from sympy.utilities.lambdify import lambdify
        from sympy.vector import CoordSys3D, express
        from sympy import symbols, pi
        from sympy import simplify, pycode
        from sympy import atan2, sqrt
        import sympy as syp
        import numpy as np

        ux, uy, uz, vx, vy, vz = symbols('ux,uy,uz,vx,vy,vz')
        # 车体坐标系为 前-左-上
        yaw = symbols('yaw')  # 假设偏航角==航向角，正北为0正东为90度，绕Zb轴转
        pitch = symbols('pitch')  # 抬头为正，绕Yb轴旋转
        roll = symbols('roll')  # 向右偏为正，绕Xb轴旋转

        ground_coord = CoordSys3D('ground')  # 东-北-天 地面坐标系

        # 通过欧拉角旋转，从东-北-天 地面坐标系 得到 UAV机体的前-右-下坐标系
        # 旋转顺序依次为：
        # Zb, Yb, Xb
        # 最开始 东-北-天（Xg-Yg-Zg)和右-前-上(Yb-Xb-(-Zb)) 是对应的
        uav_coord = ground_coord.orient_new_body(
            'uav', pi/2-yaw, pitch, roll, 'ZYX')

        # 相对风
        rel_wind_uav_coord = uy*uav_coord.i + ux*uav_coord.j + \
            uz*uav_coord.k  # 注意uy是在Xb轴上的，ux是在Yb轴上的

        # 无人机地速，
        vy = 0  # 假设vy等于0
        ground_speed_uav_coord = vx * uav_coord.i + vy*uav_coord.j + vz*uav_coord.k

        abs_wind_uav_coord = rel_wind_uav_coord + ground_speed_uav_coord

        # 把绝对风矢量按地面坐标系表示。
        ex = express(abs_wind_uav_coord, ground_coord)

        # 风速
        wX = ex.components[ground_coord.i]
        wY = ex.components[ground_coord.j]
        wZ = ex.components[ground_coord.k]
        wS = sqrt(wX**2+wY**2)
        # 风向，
        wD = syp.Mod(atan2(-wX, -wY), 2*pi)*180/pi

        self.calc_wind_speed = lambdify([ux, uy, uz, vx, vz, pitch, roll, yaw], [
                          wX, wY, wZ, wS, wD])

    def start(self):
        self.running_thread.start()

    def msg_decoder(self, msg):
        """解析udp的message"""
        assert len(msg) == struct.calcsize(self.udp_msg_format)+8
        msg = msg[6:6+struct.calcsize(self.udp_msg_format)]
        udp_values = struct.unpack(self.udp_msg_format, msg)
        physical_values = [
            uv/s+o for uv, s, o in zip(udp_values, self.udp_scalings, self.udp_offsets)]
        return physical_values

    def daq_func(self):
        while(not self.event.is_set()):
            udp_msg, (source_ip, source_port) = self.client_socket.recvfrom(512)
            if source_port != self.source_port:
                logger.debug(
                    f'skip udp msg {udp_msg} from {(source_ip,source_port)}')
                continue
            ts = time.time()
            rfc3339_ts = datetime.datetime.fromtimestamp(
                ts).astimezone(LOCAL_TIMEZONE).isoformat()
            logger.debug(f'udp_msg={udp_msg}')
            try:
                physical_values = self.msg_decoder(udp_msg)
                telemetry = {
                    'ts': int(ts*1000),
                    'values': dict(zip(self.udp_names, physical_values))
                }

                self.vx = telemetry['values']['vx']
                self.vz = telemetry['values']['vz']
                self.roll = telemetry['values']['roll']
                self.pitch = telemetry['values']['pitch']
                self.yaw = telemetry['values']['yaw']

                ux = self.ser_daq_obj.ux if self.ser_daq_obj is not None else 0
                uy = self.ser_daq_obj.uy if self.ser_daq_obj is not None else 0
                uz = self.ser_daq_obj.uz if self.ser_daq_obj is not None else 0

                wX, wY, wZ, wS, wD = self.calc_wind_speed(
                    ux,
                    uy,
                    uz,
                    self.vx,
                    self.vz,
                    self.pitch,
                    self.roll,
                    self.yaw,
                    )
                
                telemetry['values'].update(dict(zip(self.wind_speed_names,[wX,wY,wZ,wS,wD])))

                logger.debug(telemetry)

                telemetry_payload = json.dumps(telemetry)
                (rc, mid) = self.mqttc.publish(
                    topic='v1/devices/me/telemetry', payload=telemetry_payload, qos=1)
                self.data_logger.info(
                    '\t'.join([rfc3339_ts, str(ts)]+[str(pv) for pv in physical_values]))

            except KeyboardInterrupt:
                logger.debug('KeyboardInterrupt')
                break
            except Exception as e:
                logger.exception(e)
                raise
            time.sleep(0.01)
        logger.info('exitting')


class SERIAL_DAQ():
    serial_names = 'ux,uy,uz,temp'.split(',')
    serial_units = 'm/s,m/s,m/s,degC'.split(',')

    def __init__(self, data_logger, mqttc, serial_config) -> None:
        """初始化"""

        # 初始化变量
        self.data_logger = data_logger
        self.mqttc = mqttc
        self.ser = serial.Serial(**serial_config)
        self.ux = 0
        self.uy = 0
        self.uz = 0

        # 写入数据文件的首行
        sio = io.StringIO()
        print('#time\ttimestamp\t'+'\t'.join(self.serial_names), file=sio)
        print('#rfc3339\tseconds\t'+'\t'.join(self.serial_units), file=sio)
        header_rows = sio.getvalue().strip()
        sio.close()
        data_logger.info(header_rows)

        # thingsboard的设备client-side attribute中写入单位信息
        units_client_attribute = {'units': dict(
            zip(self.serial_names, self.serial_units))}
        logger.info(
            f' serial units_client_attribute = {units_client_attribute}')
        payload = json.dumps(units_client_attribute)
        (rc, mid) = mqttc.publish(
            topic='v1/devices/me/attributes',
            payload=payload,
            qos=1)

        # 数采线程控制
        self.event = threading.Event()
        self.running_thread = threading.Thread(
            target=self.daq_func, daemon=True, name='serial_daq')

    def start(self):
        self.running_thread.start()

    def msg_decoder(self, msg):
        """ 
        example: 
        """
        physical_values = [float(s) for s in msg.strip().split()]
        return physical_values

    def daq_func(self):
        while(not self.event.is_set()):
            udp_msg = self.ser.read_until(b'\n')
            ts = time.time()
            rfc3339_ts = datetime.datetime.fromtimestamp(
                ts).astimezone(LOCAL_TIMEZONE).isoformat()
            logger.debug(f'ser_msg={udp_msg}')
            try:
                physical_values = self.msg_decoder(udp_msg)
                telemetry = {
                    'ts': int(ts*1000),
                    'values': dict(zip(self.serial_names, physical_values))
                }
                telemetry_payload = json.dumps(telemetry)
                logger.debug(telemetry)
                (rc, mid) = self.mqttc.publish(
                    topic='v1/devices/me/telemetry',
                    payload=telemetry_payload,
                    qos=1)  # 单位信息发布
                self.data_logger.info(
                    '\t'.join([rfc3339_ts, str(ts)]+[str(pv) for pv in physical_values]))

                self.ux = telemetry['values']['ux']
                self.uy = telemetry['values']['uy']
                self.uz = telemetry['values']['uz']
            except KeyboardInterrupt:
                logger.debug('KeyboardInterrupt')
                break
            except Exception as e:
                logger.exception(e)
                raise
            time.sleep(0.01)
        logger.info('exitting')


def main(main_config):
    # mqtt
    token = main_config['mqtt']['token']
    host = main_config['mqtt']['host']
    port = main_config['mqtt']['port']
    # mqtt
    # doc:https://www.eclipse.org/paho/clients/python/docs/

    def on_connect(mqttc, obj, flags, rc):
        logger.info('MQTT Connected with result code '+str(rc))
        if rc != 0:
            raise ConnectionError('connection to thingsboard failed')

    mqttc_udp = mqtt.Client()
    mqttc_udp.username_pw_set(username=token['udp'])
    mqttc_udp.connect(host=host, port=port)
    mqttc_udp.on_connect = on_connect
    mqttc_udp.loop_start()  # paho自带自动重连！

    mqttc_serial = mqtt.Client()
    mqttc_serial.username_pw_set(username=token['serial'])
    mqttc_serial.connect(host=host, port=port)
    mqttc_serial.on_connect = on_connect
    mqttc_serial.loop_start()  # paho自带自动重连！

    if main_config['use_serial_port']:
        ser_daq_obj = SERIAL_DAQ(
            serial_daq,
            mqttc_serial,
            main_config['serial_port'])
    else:
        ser_daq_obj=None

    udp_daq_obj = UDP_DAQ(
        udp_daq,
        mqttc_udp,
        udp_ip=main_config['udp']['udp_ip'],
        udp_port=main_config['udp']['udp_port'],
        source_port=main_config['udp']['source_port'],
        ser_daq_obj = ser_daq_obj)

    # 加入单位

    #assert rc == mqtt.MQTT_ERR_SUCCESS

    if main_config.get('use_serial_port',True):
        ser_daq_obj.start()  # ser要先启动
    udp_daq_obj.start()

    try:
        while(True):
            time.sleep(1)
            # 计算绝对风速和风向

    except Exception as e:
        logger.exception(e)
    finally:
        udp_daq_obj.event.set()
        ser_daq_obj.event.set()
        udp_daq_obj.running_thread.join()
        ser_daq_obj.running_thread.join()


if __name__ == "__main__":
    try:
        main(config['main'])
    except Exception as e:
        logger.exception(e)
    finally:
        time.sleep(0.5)
        input('点击任意键结束')
